﻿using System;
using ePipe;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks.Sources;

namespace Letter.Net.Sockets;

class SocketAwaitableEventArgs : SocketAsyncEventArgs, IValueTaskSource<SocketOperationResult>
{
    private static readonly Action<object> _continuationCompleted = _ => { };

    private readonly PipeScheduler _ioScheduler;

    private Action<object> _continuation;

    public SocketAwaitableEventArgs(PipeScheduler ioScheduler)
#if NETCOREAPP
        : base(unsafeSuppressExecutionContextFlow: true)
#endif
    {
        _ioScheduler = ioScheduler;
    }

    protected override void OnCompleted(SocketAsyncEventArgs _)
    {
        var c = _continuation;

        if (c != null || (c = Interlocked.CompareExchange(ref _continuation, _continuationCompleted, null)) != null)
        {
            var continuationState = UserToken;
            UserToken = null;
            _continuation = _continuationCompleted; // in case someone's polling IsCompleted

            _ioScheduler.Schedule(c, continuationState);
        }
    }

    public SocketOperationResult GetResult(short token)
    {
        _continuation = null;

        if (SocketError != SocketError.Success)
        {
            return new SocketOperationResult(CreateException(SocketError));
        }

        return new SocketOperationResult(BytesTransferred);
    }

    protected static SocketException CreateException(SocketError e)
    {
        return new SocketException((int) e);
    }

    public ValueTaskSourceStatus GetStatus(short token)
    {
        return !ReferenceEquals(_continuation, _continuationCompleted) ? ValueTaskSourceStatus.Pending :
            SocketError == SocketError.Success ? ValueTaskSourceStatus.Succeeded :
            ValueTaskSourceStatus.Faulted;
    }

    public void OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags)
    {
        UserToken = state;
        var prevContinuation = Interlocked.CompareExchange(ref _continuation, continuation, null);
        if (ReferenceEquals(prevContinuation, _continuationCompleted))
        {
            UserToken = null;

#if NETSTANDARD2_1
            ThreadPool.QueueUserWorkItem((o =>
            {
                var item = (Tuple<Action<object>, object>) o;
                var action = item.Item1;
                var args = item.Item2;
                action(args);
            }),
            Tuple.Create<Action<object>, object>(continuation, state));
#elif NETCOREAPP
            ThreadPool.UnsafeQueueUserWorkItem(continuation, state, preferLocal: true);
#endif
        }
    }
}